from PIL import Image
import numpy as np
import os

def create_dummy_data(base_dir="./data", num_samples=10, image_size=(256, 256), num_classes=2):
    """创建假图像和标签数据"""
    
    image_dir = os.path.join(base_dir, "images")
    label_dir = os.path.join(base_dir, "labels")
    inference_dir = os.path.join(base_dir, "inference_images")
    
    os.makedirs(image_dir, exist_ok=True)
    os.makedirs(label_dir, exist_ok=True)
    os.makedirs(inference_dir, exist_ok=True)
    
    train_list = []
    val_list = []
    test_list = []
    
    print(f"开始生成 {num_samples} 个假数据样本...")
    for i in range(num_samples):
        # 创建假图像 (RGB)
        img_array = np.random.randint(0, 256, (image_size[0], image_size[1], 3), dtype=np.uint8)
        img = Image.fromarray(img_array)
        img_filename = f"image_{i:03d}.png"
        img.save(os.path.join(image_dir, img_filename))
        
        # 创建假标签 (灰度，像素值对应类别ID)
        mask_array = np.random.randint(0, num_classes, (image_size[0], image_size[1]), dtype=np.uint8)
        mask = Image.fromarray(mask_array)
        mask_filename = f"image_{i:03d}.png"
        mask.save(os.path.join(label_dir, mask_filename))
        
        # 划分数据集
        if i < num_samples * 0.7: # 70% 训练
            train_list.append(f"image_{i:03d}") # 不带后缀
        elif i < num_samples * 0.85: # 15% 验证
            val_list.append(f"image_{i:03d}") # 不带后缀
        else: # 15% 测试
            test_list.append(f"image_{i:03d}") # 不带后缀
            
        # 复制一些图像到推理目录
        if i % 3 == 0: # 每隔3个复制一个用于推理测试
            img.save(os.path.join(inference_dir, img_filename))
            
    # 写入文件列表
    with open(os.path.join(base_dir, "train.txt"), "w") as f:
        for fname in train_list:
            f.write(fname + "\n")
            
    with open(os.path.join(base_dir, "val.txt"), "w") as f:
        for fname in val_list:
            f.write(fname + "\n")
            
    with open(os.path.join(base_dir, "test.txt"), "w") as f:
        for fname in test_list:
            f.write(fname + "\n")
            
    print("假数据生成完成！")
    print(f"图像保存在: {image_dir}")
    print(f"标签保存在: {label_dir}")
    print(f"推理图像保存在: {inference_dir}")
    print(f"train.txt: {len(train_list)} 个样本")
    print(f"val.txt: {len(val_list)} 个样本")
    print(f"test.txt: {len(test_list)} 个样本")

if __name__ == "__main__":
    # 在项目根目录下运行此脚本
    # 确保当前工作目录是 medical_segmentation_project
    create_dummy_data(base_dir="./data", num_samples=20, image_size=(128, 128), num_classes=2)


